# This script creates a connector for the BEG embedding model in OpenSearch Service.
# It uses the AWS4Auth library to sign the request with AWS credentials.

import boto3
import json
import os
import requests 
from requests_aws4auth import AWS4Auth

opensearch_service_api_endpoint = os.environ['OPENSEARCH_SERVICE_DOMAIN_ENDPOINT']
region = os.environ['DEEPSEEK_AWS_REGION']
invoke_role_arn = os.environ['INVOKE_DEEPSEEK_ROLE']
create_deepseek_connector_role_arn = os.environ['CREATE_DEEPSEEK_CONNECTOR_ROLE']
sagemaker_endpoint_url = os.environ['SAGEMAKER_EMBEDDING_MODEL_INFERENCE_ENDPOINT']


# Create the AWS4Auth object that will sign the create connector API call. 
credentials = boto3.client('sts').assume_role(
    RoleArn=create_deepseek_connector_role_arn,
    RoleSessionName='create_connector_session'
)['Credentials']
awsauth = AWS4Auth(credentials['AccessKeyId'], 
                   credentials['SecretAccessKey'], 
                   region, 
                   'es', 
                   session_token=credentials['SessionToken'])


# Prepare the API call parameters. 
path = '/_plugins/_ml/connectors/_create'
url = opensearch_service_api_endpoint + path
# See the documentation 
# https://opensearch.org/docs/latest/ml-commons-plugin/remote-models/blueprints/ for 
# details on the connector payload, and additional blueprints for other models.
payload = {
  "name": "BEG Embedding model connector",
  "description": "Connector for Sagemaker BEG model",
  "version": "1.0",
  "protocol": "aws_sigv4",
  "credential": {
    "roleArn": invoke_role_arn
  },
  "parameters": {
    "service_name": "sagemaker",
    "region": region
  },
    "actions": [
        {
            "action_type": "PREDICT",
            "method": "POST",
            "url": sagemaker_endpoint_url,
            "headers": {
                "content-type": "application/json"
            },
            "request_body": "{ \"input\": \"${parameters.input}\"}",
            "pre_process_function": """
                StringBuilder builder = new StringBuilder();
                builder.append("\\"");
                builder.append(params.text_docs[0]);
                builder.append("\\"");
                def parameters = "{" +"\\"input\\":" + builder + "}";
                return "{" +"\\"parameters\\":" + parameters + "}";
                """,
            "post_process_function": """
                  def name = "sentence_embedding";
                  def dataType = "FLOAT32";
                  if (params.data[0].embedding == null || params.data[0].embedding.length == 0) {
                    return params.message;
                  }
                  def shape = [params.data[0].embedding.length];
                  def json = "{" +
                             "\\"name\\":\\"" + name + "\\"," +
                             "\\"data_type\\":\\"" + dataType + "\\"," +
                             "\\"shape\\":" + shape + "," +
                             "\\"data\\":" + params.data[0].embedding +
                             "}";
                  return json;
                """
        }
    ]
}

# This ignores errors and doesn't check the result. In real use,
# you should wrap this code with try/except blocks and check the
# response status code and the response body for errors.
headers = {"Content-Type": "application/json"}
r = requests.post(url, auth=awsauth, json=payload, headers=headers)
print(r.text)
connector_id = json.loads(r.text)['connector_id']


print(connector_id)
print(f'EMBEDDING_CONNECTOR_ID="{connector_id}"\n')
os.environ["EMBEDDING_CONNECTOR_ID"] = connector_id
